{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yeadDkMiISin"
      },
      "source": [
        "# PaLM API: Text Quickstart with Python"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lEXQ3OwKIa-O"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/palm_docs/text_quickstart\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on ai.google.dev</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/palm_docs/text_quickstart.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/generative-ai-docs/blob/main/site/en/palm_docs/text_quickstart.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jp_CKyzxUqx6"
      },
      "source": [
        "In this notebook, you'll learn how to get started with the PaLM API, which gives you access to Google's latest large language models. Here, you'll learn how to use the PaLM API's text generation features."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gQnbmB70zqon"
      },
      "source": [
        "## Setup\n",
        "\n",
        "**Note**: At this time, the PaLM API is [only available in certain regions](https://developers.generativeai.google/available_regions)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gePcZFDdypdP"
      },
      "source": [
        "First, download and install the PaLM API Python library."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "Mk2d90cCdF4u"
      },
      "outputs": [],
      "source": [
        "!pip install -q google-generativeai"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "zCm_GdCKYDNm"
      },
      "outputs": [],
      "source": [
        "import pprint\n",
        "import google.generativeai as palm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taFqmFg-7_PY"
      },
      "source": [
        "### Grab an API Key\n",
        "\n",
        "To get started, you'll need to [create an API key](https://developers.generativeai.google/tutorials/setup)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "YE1x5qv-hka3"
      },
      "outputs": [],
      "source": [
        "palm.configure(api_key='YOUR_API_KEY')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G2dk6P3nJz6m"
      },
      "source": [
        "## Text generation\n",
        "\n",
        "Use the `palm.list_models` function to find available models:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "dKLCxpTw_C2Q"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "models/text-bison-001\n"
          ]
        }
      ],
      "source": [
        "models = [m for m in palm.list_models() if 'generateText' in m.supported_generation_methods]\n",
        "model = models[0].name\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9OvvPDyVvb4X"
      },
      "source": [
        "Use the `palm.generate_text` method to generate text:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "sJm90x7J6YLe"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "There are 3 houses * 3 cats / house = 9 cats. So, 9 cats * 4 mittens / cat = 36 mittens were made. Also, 9 cats * 1 hat / cat = 9 hats were made. So, 36 mittens * 7m / mitten = 252m of yarn was used for the mittens. Also, 9 hats * 4m / hat = 36m of yarn was used for the hats. In total, 252m + 36m = 288m of yarn was used.\n",
            "Thus, the answer is 288.\n"
          ]
        }
      ],
      "source": [
        "prompt = \"\"\"\n",
        "You are an expert at solving word problems.\n",
        "\n",
        "Solve the following problem:\n",
        "\n",
        "I have three houses, each with three cats.\n",
        "each cat owns 4 mittens, and a hat. Each mitten was\n",
        "knit from 7m of yarn, each hat from 4m.\n",
        "How much yarn was needed to make all the items?\n",
        "\n",
        "Think about it step by step, and show your work.\n",
        "\"\"\"\n",
        "\n",
        "completion = palm.generate_text(\n",
        "    model=model,\n",
        "    prompt=prompt,\n",
        "    temperature=0,\n",
        "    # The maximum length of the response\n",
        "    max_output_tokens=800,\n",
        ")\n",
        "\n",
        "print(completion.result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5eca9b28591a"
      },
      "source": [
        "## More options\n",
        "\n",
        "The `palm.generate_text` function has a few other arguments worth mentioning."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dede032595a0"
      },
      "source": [
        "### Stop sequences\n",
        "\n",
        "Use the `stop_sequences` argument to stop generation early. \n",
        "\n",
        "For example LLM's often make mistakes in arithmetic. You could ask the model to \"use a calculator\" by putting equations in a `<calc>` tag. \n",
        "\n",
        "Have the model stop at the closing tag, so you can edit the prompt:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "ecc8e5092861"
      },
      "outputs": [],
      "source": [
        "calc_prompt = f\"\"\"\n",
        "Please solve the following problem.\n",
        "\n",
        "{prompt}\n",
        "\n",
        "----------------\n",
        "\n",
        "Important: Use the calculator for each step.\n",
        "Don't do the arithmetic in your head. \n",
        "\n",
        "To use the calculator wrap an equation in <calc> tags like this: \n",
        "\n",
        "<calc> 3 cats * 2 hats/cat </calc> = 6\n",
        "\n",
        "----------------\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "equation=None\n",
        "while equation is None:\n",
        "    completion = palm.generate_text(\n",
        "        model=model,\n",
        "        prompt=calc_prompt,\n",
        "        stop_sequences=['</calc>'],\n",
        "        # The maximum length of the response\n",
        "        max_output_tokens=800,\n",
        "    )\n",
        "\n",
        "    try:\n",
        "        response, equation = completion.result.split('<calc>', maxsplit=1)\n",
        "    except Exception:\n",
        "        continue"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "21cd8789a136"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Chain-of-thought:\n",
            "\n",
            "There are three houses, and each house has three cats, so there are 3 houses * 3 cats / house = 9 cats. Each cat has 4 mittens, so the cats need 9 cats * 4 mittens / cat = 36 mittens. Each mitten takes 7m of yarn, so 36 mittens * 7m / mitten = 252m of yarn. Each cat has a hat, and each hat takes 4m of yarn, so 9 cats * 4m / cat = 36m of yarn. So, in total, 36m + 252m = 288m of yarn were needed.\n",
            "\n",
            "The answer should be \n"
          ]
        }
      ],
      "source": [
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "6f23f57698f2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " 9 cats * 4 mittens / cat \n"
          ]
        }
      ],
      "source": [
        "print(equation)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "13eda23afe44"
      },
      "source": [
        "From there you can calculate the result, and assemble a new prompt for the model to continue from.\n",
        "For a complete working implementation see the [Text calculator example](../examples/text_calculator.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d1b0509ef19e"
      },
      "source": [
        "### Candidates\n",
        "\n",
        "Typically, there's some degree of randomness in the text produced by LLMs. (Read more about why in the [LLM primer](https://developers.generativeai.google/guide/concepts)). That means that when you call the API more than once with the same input, you might get different responses. You can use this feature to your advantage to get alternate model responses.\n",
        "\n",
        "The `temperature` argument controls the variance of the responses. The `palm.Model` object gives the default value for `temperature` and other parameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "s6doS7vAvxz8"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Model(name='models/text-bison-001', base_model_id='', version='001', display_name='Text Bison', description='Model targeted for text generation.', input_token_limit=8196, output_token_limit=1024, supported_generation_methods=['generateText'], temperature=0.7, top_p=0.95, top_k=40)"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "models[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GYJbWwzXwXiD"
      },
      "source": [
        "The `candidate_count` argument controls the number of responses returned:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "NCvc78t-ROrL"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "In each house there are 3 cats * 4 mittens / cat = 12 mittens. In total there are 3 houses * 12 mittens / house = 36 mittens. In total there are 36 mittens * 7m / mitten = 252m of yarn for the mittens. In total there are 3 houses * 3 cats / house * 1 hat / cat = 9 hats. In total there are 9 hats * 4m / hat = 36m of yarn for the hats. In total there are 36m yarn for the hats + 252m yarn for the mittens = 288m of yarn.\n",
            "The answer: 288.\n"
          ]
        }
      ],
      "source": [
        "completion = palm.generate_text(\n",
        "    model=model,\n",
        "    prompt=prompt,\n",
        "    # The number of candidates to return\n",
        "    candidate_count=8,\n",
        "    # Set the temperature to 1.0 for more variety of responses.\n",
        "    temperature=1.0,\n",
        "    max_output_tokens=800,\n",
        ")\n",
        "\n",
        "print(completion.result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2lXu8iu0Mm3V"
      },
      "source": [
        "When you request multiple candidates the `Completion.result` attribute still just contains the first one. The `Completion.candidates` attribute contains all of them:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "nJEjPHGVMPaj"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[{'output': 'In each house there are 3 cats * 4 mittens / cat = 12 mittens. In '\n",
            "            'total there are 3 houses * 12 mittens / house = 36 mittens. In '\n",
            "            'total there are 36 mittens * 7m / mitten = 252m of yarn for the '\n",
            "            'mittens. In total there are 3 houses * 3 cats / house * 1 hat / '\n",
            "            'cat = 9 hats. In total there are 9 hats * 4m / hat = 36m of yarn '\n",
            "            'for the hats. In total there are 36m yarn for the hats + 252m '\n",
            "            'yarn for the mittens = 288m of yarn.\\n'\n",
            "            'The answer: 288.',\n",
            "  'safety_ratings': [{'category': <HarmCategory.HARM_CATEGORY_DEROGATORY: 1>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_TOXICITY: 2>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_VIOLENCE: 3>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_SEXUAL: 4>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_MEDICAL: 5>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_DANGEROUS: 6>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>}]},\n",
            " {'output': 'Each house has 3 cats, so each house needs 3 * 4 = 12 mittens. '\n",
            "            \"With three houses, that's 3 * 12 = 36 mittens. And each house \"\n",
            "            'needs 3 * 1 = 3 hats. So in total, we need 3 hats + 36 mittens = '\n",
            "            '39 items. Each mitten needs 7 meters of yarn, so 39 mittens need '\n",
            "            '39 * 7 = 273 meters of yarn. Each hat needs 4 meters of yarn, and '\n",
            "            \"we need 3 hats, so that's 4 * 3 = 12 meters of yarn. So in total, \"\n",
            "            'we needed 12 + 273 = 285 meters of yarn.\\n'\n",
            "            'Thus, the answer is 285.',\n",
            "  'safety_ratings': [{'category': <HarmCategory.HARM_CATEGORY_DEROGATORY: 1>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_TOXICITY: 2>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_VIOLENCE: 3>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_SEXUAL: 4>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_MEDICAL: 5>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_DANGEROUS: 6>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>}]},\n",
            " {'output': 'There are 3 houses * 3 cats / house = 9 cats. There are 9 cats * '\n",
            "            '4 mittens / cat = 36 mittens. There are 9 cats * 1 hat / cat = 9 '\n",
            "            'hats. The total amount of yarn for the mittens is 36 mittens * 7m '\n",
            "            '/ mitten = 252m. The total amount of yarn for the hats is 9 hats '\n",
            "            '* 4m / hat = 36m. The total amount of yarn is 252m + 36m = 288m.\\n'\n",
            "            'Thus, the answer is 288.',\n",
            "  'safety_ratings': [{'category': <HarmCategory.HARM_CATEGORY_DEROGATORY: 1>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_TOXICITY: 2>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_VIOLENCE: 3>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_SEXUAL: 4>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_MEDICAL: 5>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_DANGEROUS: 6>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>}]},\n",
            " {'output': 'There are 3 houses * 3 cats / house = 9 cats. Each cat has 4 '\n",
            "            'mittens + 1 hat = 5 items. So the total number of items is 9 cats '\n",
            "            '* 5 items / cat = 45 items. Thus, 45 items * 7m / item = 315m of '\n",
            "            'yarn was needed.\\n'\n",
            "            'Thus, the answer is 315.',\n",
            "  'safety_ratings': [{'category': <HarmCategory.HARM_CATEGORY_DEROGATORY: 1>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_TOXICITY: 2>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_VIOLENCE: 3>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_SEXUAL: 4>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_MEDICAL: 5>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_DANGEROUS: 6>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>}]},\n",
            " {'output': 'Chain-of-thought:\\n'\n",
            "            'There are 3 houses * 3 cats / house = 9 cats.\\n'\n",
            "            'The cats need 9 cats * 4 mittens / cat = 36 mittens.\\n'\n",
            "            'The cats need 9 cats * 1 hat / cat = 9 hats.\\n'\n",
            "            'The mittens need 36 mittens * 7m / mitten = 252m of yarn.\\n'\n",
            "            'The hats need 9 hats * 4m / hat = 36m of yarn.\\n'\n",
            "            'Therefore, the total amount of yarn needed is 252m + 36m = 288m.\\n'\n",
            "            '\\n'\n",
            "            'The answer should be 288',\n",
            "  'safety_ratings': [{'category': <HarmCategory.HARM_CATEGORY_DEROGATORY: 1>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_TOXICITY: 2>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_VIOLENCE: 3>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_SEXUAL: 4>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_MEDICAL: 5>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_DANGEROUS: 6>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>}]},\n",
            " {'output': 'First find the total number of cats: 3 houses * 3 cats / house = '\n",
            "            '9 cats. Then multiply that number by the number of mittens per '\n",
            "            'cat to find the total number of mittens: 9 cats * 4 mittens / cat '\n",
            "            '= 36 mittens. Then multiply that number by the number of meters '\n",
            "            'of yarn per mitten to find the total amount of yarn used for '\n",
            "            'mittens: 36 mittens * 7 meters / mitten = 252 meters. Then do the '\n",
            "            'same thing for hats: 9 cats * 1 hat / cat = 9 hats. Then multiply '\n",
            "            'that number by the number of meters of yarn per hat to find the '\n",
            "            'total amount of yarn used for hats: 9 hats * 4 meters / hat = 36 '\n",
            "            'meters. Then add the amount of yarn used for mittens and hats to '\n",
            "            'find the total amount of yarn used: 36 meters + 252 meters = 288 '\n",
            "            'meters.\\n'\n",
            "            'Thus, the answer is 288.',\n",
            "  'safety_ratings': [{'category': <HarmCategory.HARM_CATEGORY_DEROGATORY: 1>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_TOXICITY: 2>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_VIOLENCE: 3>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_SEXUAL: 4>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_MEDICAL: 5>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_DANGEROUS: 6>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>}]},\n",
            " {'output': 'The total number of cats is 3 houses * 3 cats / house = 9 cats. '\n",
            "            'So, the total number of mittens is 9 cats * 4 mittens / cat = 36 '\n",
            "            'mittens. The total number of hats is 9 cats * 1 hat / cat = 9 '\n",
            "            'hats. The total length of yarn needed to make the mittens is 36 '\n",
            "            'mittens * 7 m / mitten = 252 m. The total length of yarn needed '\n",
            "            'to make the hats is 9 hats * 4 m / hat = 36 m. So, the total '\n",
            "            'length of yarn needed is 252 m + 36 m = 288 m.\\n'\n",
            "            '\\n'\n",
            "            'The answer: 288',\n",
            "  'safety_ratings': [{'category': <HarmCategory.HARM_CATEGORY_DEROGATORY: 1>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_TOXICITY: 2>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_VIOLENCE: 3>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_SEXUAL: 4>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_MEDICAL: 5>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_DANGEROUS: 6>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>}]},\n",
            " {'output': 'There are 3 houses with 3 cats each, so 3 * 3 = 9 cats. Each cat '\n",
            "            'has 4 mittens and a hat, so 9 cats * 4 mittens / cat + 9 cats * 1 '\n",
            "            'hat / cat = 36 mittens and 9 hats. Each mitten takes 7m of yarn '\n",
            "            'and each hat takes 4m of yarn, so the total yarn needed is 36 '\n",
            "            'mittens * 7m / mitten + 9 hats * 4m / hat = 252m + 36m = 288m.\\n'\n",
            "            'The answer: 288.',\n",
            "  'safety_ratings': [{'category': <HarmCategory.HARM_CATEGORY_DEROGATORY: 1>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_TOXICITY: 2>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_VIOLENCE: 3>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_SEXUAL: 4>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_MEDICAL: 5>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>},\n",
            "                     {'category': <HarmCategory.HARM_CATEGORY_DANGEROUS: 6>,\n",
            "                      'probability': <HarmProbability.NEGLIGIBLE: 1>}]}]\n"
          ]
        }
      ],
      "source": [
        "import pprint\n",
        "pprint.pprint(completion.candidates)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b5ad0836a2f6"
      },
      "source": [
        "So, since you know the answer to this problem, it's easy to check the solve rate:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "xJh2LRHhxQ6k"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0.75"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import numpy as np\n",
        "np.mean(['288' in c['output'] for c in completion.candidates])"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "text_quickstart.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
